import heapq
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import defaultdict, Counter
import networkx as nx

# 树节点类
class Node:
    def __init__(self, frequency, symbol, left=None, right=None):
        self.frequency = frequency
        self.symbol = symbol
        self.left = left
        self.right = right
        self.huffman_code = ''

    # 定义比较方法以使节点可以放入优先队列
    def __lt__(self, other):
        return self.frequency < other.frequency

# 生成哈夫曼树
def huffman_tree(data):
    frequency = Counter(data)
    priority_queue = [Node(frequency=symbol[1], symbol=symbol[0]) for symbol in frequency.items()]
    heapq.heapify(priority_queue)

    frames = []
    G = nx.DiGraph()

    while len(priority_queue) > 1:
        left = heapq.heappop(priority_queue)
        right = heapq.heappop(priority_queue)

        merged = Node(left.frequency + right.frequency, left.symbol + right.symbol, left, right)
        heapq.heappush(priority_queue, merged)

        G.add_edge(left.symbol, merged.symbol)
        G.add_edge(right.symbol, merged.symbol)
        frames.append(G.copy())

    return priority_queue[0], frames

# 绘制哈夫曼树
def draw_huffman_tree(frames, interval=1000):
    fig, ax = plt.subplots(figsize=(8, 6))

    def update(num):
        ax.clear()
        G = frames[num]
        pos = nx.spring_layout(G)
        nx.draw(G, pos, with_labels=True, node_size=3000, node_color="lightblue", font_size=10, font_weight="bold", ax=ax)
        ax.set_title(f"Step {num + 1}")

    ani = animation.FuncAnimation(fig, update, frames=len(frames), interval=interval, repeat=False)
    plt.show()

# 示例数据
data = [1, 2, 3, 4, 5, 1, 2, 3, 56]

# 生成哈夫曼树并获取动画帧
root, frames = huffman_tree(data)

# 绘制哈夫曼树生成过程的动画
draw_huffman_tree(frames)
